import numpy as np
import matplotlib.pyplot as plt

# 设置参数
c = 1500  # 声速
fs = 8000  # 采样频率
f = 1000  # 信号频率
d = 0.2  # 水听器间距
SNR = 10  # 信噪比
num_trials = 10  # 实验次数


def generate_signal_mean_intensity(theta, num_samples):
    """
    生成指定方位的信号
    """
    print(f"生成的信号方向为：{np.rad2deg(theta)}")
    t = np.arange(num_samples) / fs  # 采样时刻序列
    pressure = np.sin(2 * np.pi * f * t)
    velocity_x = np.cos(theta) * pressure  # x轴振速与声压同相
    velocity_y = np.sin(theta) * pressure  # y轴振速与声压同相
    noise_power = 10 ** (-SNR / 10)
    noise_p = np.random.randn(num_samples) * np.sqrt(noise_power)
    noise_nx = np.random.randn(num_samples) * np.sqrt(noise_power/2)
    noise_ny = np.random.randn(num_samples) * np.sqrt(noise_power / 2)

    return pressure + noise_p, velocity_x + noise_nx, velocity_y + noise_ny


def estimate_doa_intensity(pressure, velocity_x, velocity_y):
    """
    平均声强法估计DOA
    """
    # 计算x轴和y轴的声强
    intensity_x = pressure * velocity_x
    intensity_y = pressure * velocity_y

    # 计算平均声强
    mean_intensity_x = np.mean(intensity_x)
    mean_intensity_y = np.mean(intensity_y)

    # 估计DOA
    theta_est = np.arctan2(mean_intensity_y, mean_intensity_x)

    return theta_est


# 仿真实验
errors = []
for _ in range(num_trials):
    # 生成随机方位角
    theta_true = np.random.uniform(-np.pi/2, np.pi/2)

    # 生成信号
    p, v_x, v_y = generate_signal_mean_intensity(theta_true, 1024)
    # signal2 = generate_signal(theta_true, 1024)

    # 估计DOA
    theta_est = estimate_doa_intensity(p, v_x, v_y)
    print(f"估计的信号方向为{np.rad2deg(theta_est)}")
    # 排除错误解算
    if theta_est is not None:
        errors.append(np.abs(theta_true - theta_est))

# 计算平均误差
mean_error = np.mean(errors) * 180 / np.pi
print("mean error: {:.2f} degree".format(mean_error))

# # 绘制误差分布图
# plt.hist(errors, bins=10)
# plt.xlabel("error (rad)")
# plt.ylabel("times")
# plt.show()